#include <cmath>
#include <vector>
#include <string>
#include "aloam_velodyne/common.h"
#include "aloam_velodyne/tic_toc.h"
#include <nav_msgs/Odometry.h>
#include <opencv2/opencv.hpp>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/filters/voxel_grid.h>
#include <pcl/kdtree/kdtree_flann.h>
#include <ros/ros.h>
#include <sensor_msgs/Imu.h>
#include <sensor_msgs/PointCloud2.h>
#include <tf/transform_datatypes.h>
#include <tf/transform_broadcaster.h>

using std::atan2;
using std::cos;
using std::sin;

const double scanPeriod = 0.1;
const int systemDelay = 5;
int systemInitCount = 0;
bool systemInited = false;
int N_SCANS = 0;
float cloudCurvature[400000];
int cloudSortInd[400000];
int cloudNeighborPicked[400000];
int cloudLabel[400000];
int BadcloudMarked[400000];
float range_vec[400000];

bool comp(int i, int j) { return (cloudCurvature[i] < cloudCurvature[j]); }

ros::Publisher pubLaserCloud;
ros::Publisher pubCornerPointsSharp;
ros::Publisher pubCornerPointsLessSharp;
ros::Publisher pubSurfPointsFlat;
ros::Publisher pubSurfPointsLessFlat;
ros::Publisher pubRemovePoints;
std::vector<ros::Publisher> pubEachScan;
bool PUB_EACH_LINE = false;
double MINIMUM_RANGE = 0.1;
double MAXIMUM_RANGE = 100.0;
Eigen::Quaterniond q_odom, q_odom_curr, q_odom_last;
Eigen::Vector3d t_odom, t_odom_curr, t_odom_last;

template<typename PointT>
void removeClosedPointCloud(const pcl::PointCloud<PointT> &cloud_in,
                            pcl::PointCloud<PointT> &cloud_out, float thres)
{
    if (&cloud_in != &cloud_out)
    {
        cloud_out.header = cloud_in.header;
        cloud_out.points.resize(cloud_in.points.size());
    }
    size_t j = 0;
    for (size_t i = 0; i < cloud_in.points.size(); ++i)
    {
        if (cloud_in.points[i].x * cloud_in.points[i].x + cloud_in.points[i].y * cloud_in.points[i].y +
            cloud_in.points[i].z * cloud_in.points[i].z < thres * thres) continue;
        cloud_out.points[j] = cloud_in.points[i];
        j++;
    }
    if (j != cloud_in.points.size()) cloud_out.points.resize(j);
    cloud_out.height = 1;
    cloud_out.width = static_cast<uint32_t>(j);
    cloud_out.is_dense = true;
}

void RemoveNaNFarClosePoints(const pcl::PointCloud<RsPointXYZIRT>::Ptr &cloud_in,
                             pcl::PointCloud<PointType>::Ptr &cloud_out)
{
    int cloudSize = cloud_in->size();
    RsPointXYZIRT p_t;
    PointType p_out;
    pcl::PointCloud<PointType>::Ptr cloud_filter(new pcl::PointCloud<PointType>());
    for (int i = 0; i < cloudSize; i++)
    {
        p_t = cloud_in->points[i];
        float dis = p_t.x * p_t.x + p_t.y * p_t.y;
        if (std::isnan(p_t.x) || std::isnan(p_t.y) || std::isnan(p_t.z)) continue;
        if (std::isinf(p_t.x) || std::isinf(p_t.y) || std::isinf(p_t.z)) continue;
        if (!std::isfinite(p_t.x) || !std::isfinite(p_t.y) || !std::isfinite(p_t.z)) continue;
        if (dis < MINIMUM_RANGE * MINIMUM_RANGE) continue;
        if (dis > MAXIMUM_RANGE * MAXIMUM_RANGE) continue;
        p_out.x = p_t.x;
        p_out.y = p_t.y;
        p_out.z = p_t.z;
        p_out.intensity = p_t.intensity;
        cloud_out->push_back(p_out);
    }
}

void RemoveLidarDistortion(const pcl::PointCloud<PointType>::Ptr &cloud_in, pcl::PointCloud<PointType>::Ptr &cloud_out,
                           const Eigen::Quaterniond &dqlc, const Eigen::Vector3d &dtlc)
{
    // 作业：
    // TODO: 完成点云运动畸变去除函数，利用里程计结果对每个点进行校正；
    // 答案：
    // start:
    int cloudSize = cloud_in->size();
    float startOri = -atan2(cloud_in->points[0].y, cloud_in->points[0].x);
    float endOri = -atan2(cloud_in->points[cloudSize - 1].y, cloud_in->points[cloudSize - 1].x) + 2 * M_PI;
    if (endOri - startOri > 3 * M_PI) endOri -= 2 * M_PI;
    else if (endOri - startOri < M_PI) endOri += 2 * M_PI;
    bool halfPassed = false;
    PointType p_in, p_out;
    Eigen::Vector3d startP, endP;
    size_t j = 0;
    for (int i = 0; i < cloudSize; i++)
    {
        p_in = cloud_in->points[i];
        float ori = -atan2(p_in.y, p_in.x);
        if (!halfPassed)
        {
            if (ori < startOri - M_PI / 2) ori += 2 * M_PI;
            else if (ori > startOri + M_PI * 3 / 2) ori -= 2 * M_PI;
            if (ori - startOri > M_PI) halfPassed = true;
        }
        else
        {
            ori += 2 * M_PI;
            if (ori < endOri - M_PI * 3 / 2) ori += 2 * M_PI;
            else if (ori > endOri + M_PI / 2) ori -= 2 * M_PI;
        }
        if (ori > endOri) ori -= 2 * M_PI;
        float ratio = (ori - startOri) / (endOri - startOri);
        float s = 1.0 - ratio;
        Eigen::Quaterniond delta_qlc = Eigen::Quaterniond::Identity().slerp(s, dqlc).normalized();
        const Eigen::Vector3d delta_Plc = s * dtlc;
        endP = delta_qlc * Eigen::Vector3d(p_in.x, p_in.y, p_in.z) + delta_Plc;
        p_out = p_in;
        p_out.x = endP.x();
        p_out.y = endP.y();
        p_out.z = endP.z();
        j++;
        cloud_out->push_back(p_out);
    }
    // end.
}

void odomHandler(const nav_msgs::Odometry::ConstPtr &odomMsg)
{
    q_odom.x() = odomMsg->pose.pose.orientation.x;
    q_odom.y() = odomMsg->pose.pose.orientation.y;
    q_odom.z() = odomMsg->pose.pose.orientation.z;
    q_odom.w() = odomMsg->pose.pose.orientation.w;
    q_odom.normalize();
    t_odom.x() = odomMsg->pose.pose.position.x;
    t_odom.y() = odomMsg->pose.pose.position.y;
    t_odom.z() = odomMsg->pose.pose.position.z;
}

void laserCloudHandler(const sensor_msgs::PointCloud2ConstPtr &laserCloudMsg)
{
    if (!systemInited)
    {
        systemInitCount++;
        if (systemInitCount >= systemDelay) systemInited = true;
        else
        {
            q_odom_last = q_odom;
            t_odom_last = t_odom;
            return;
        }
    }
    q_odom_curr = q_odom;
    t_odom_curr = t_odom;
    TicToc t_whole;
    TicToc t_prepare;
    std::vector<int> scanStartInd(N_SCANS, 0);
    std::vector<int> scanEndInd(N_SCANS, 0);
    pcl::PointCloud<RsPointXYZIRT>::Ptr laserCloudIn(new pcl::PointCloud<RsPointXYZIRT>());
    pcl::fromROSMsg(*laserCloudMsg, *laserCloudIn);
    pcl::PointCloud<PointType>::Ptr laserCloudFilter(new pcl::PointCloud<PointType>());
    RemoveNaNFarClosePoints(laserCloudIn, laserCloudFilter);
    Eigen::Quaterniond dq_odom;
    Eigen::Vector3d dt_odom;
    dq_odom = q_odom_curr.conjugate() * q_odom_last;
    dt_odom = q_odom_curr.conjugate() * (t_odom_last - t_odom_curr);
    dq_odom.normalize();
    pcl::PointCloud<PointType>::Ptr laserCloudUndis(new pcl::PointCloud<PointType>());
    // 作业：
    // TODO: 完成点云运动畸变去除函数，利用里程计结果对每个点进行校正；
    RemoveLidarDistortion(laserCloudFilter, laserCloudUndis, dq_odom, dt_odom);
    int cloudSize = laserCloudUndis->points.size();
    int count = cloudSize;
    PointType point;
    std::vector<pcl::PointCloud<PointType>> laserCloudScans(N_SCANS);
    for (int i = 0; i < cloudSize; i++)
    {
        point.x = laserCloudUndis->points[i].x;
        point.y = laserCloudUndis->points[i].y;
        point.z = laserCloudUndis->points[i].z;
        float angle = atan(point.z / sqrt(point.x * point.x + point.y * point.y)) * 180 / M_PI;
        int scanID = int((angle + 15) / 2 + 0.5);
        if (scanID > (N_SCANS - 1) || scanID < 0)
        {
            count--;
            continue;
        }
        point.intensity = scanID;
        laserCloudScans[scanID].push_back(point);
    }
    cloudSize = count;
    pcl::PointCloud<PointType>::Ptr laserCloud(new pcl::PointCloud<PointType>());
    for (int i = 0; i < N_SCANS; i++)
    {
        scanStartInd[i] = laserCloud->size() + 5;
        *laserCloud += laserCloudScans[i];
        scanEndInd[i] = laserCloud->size() - 6;
    }
    for (int i = 5; i < cloudSize - 5; i++)
    {
        float diffX = laserCloud->points[i - 5].x + laserCloud->points[i - 4].x + laserCloud->points[i - 3].x +
                      laserCloud->points[i - 2].x + laserCloud->points[i - 1].x - 10 * laserCloud->points[i].x +
                      laserCloud->points[i + 1].x + laserCloud->points[i + 2].x + laserCloud->points[i + 3].x +
                      laserCloud->points[i + 4].x + laserCloud->points[i + 5].x;
        float diffY = laserCloud->points[i - 5].y + laserCloud->points[i - 4].y + laserCloud->points[i - 3].y +
                      laserCloud->points[i - 2].y + laserCloud->points[i - 1].y - 10 * laserCloud->points[i].y +
                      laserCloud->points[i + 1].y + laserCloud->points[i + 2].y + laserCloud->points[i + 3].y +
                      laserCloud->points[i + 4].y + laserCloud->points[i + 5].y;
        float diffZ = laserCloud->points[i - 5].z + laserCloud->points[i - 4].z + laserCloud->points[i - 3].z +
                      laserCloud->points[i - 2].z + laserCloud->points[i - 1].z - 10 * laserCloud->points[i].z +
                      laserCloud->points[i + 1].z + laserCloud->points[i + 2].z + laserCloud->points[i + 3].z +
                      laserCloud->points[i + 4].z + laserCloud->points[i + 5].z;
        cloudCurvature[i] = diffX * diffX + diffY * diffY + diffZ * diffZ;
        cloudSortInd[i] = i;
        cloudNeighborPicked[i] = 0;
        cloudLabel[i] = 0;
        BadcloudMarked[i] = 0;
    }
    // 计算点深度；
    for (int i = 0; i < cloudSize; i++)
    {
        range_vec[i] = sqrt(
                laserCloud->points[i].x * laserCloud->points[i].x + laserCloud->points[i].y * laserCloud->points[i].y +
                laserCloud->points[i].z * laserCloud->points[i].z);  //
    }
    for (int i = 5; i < cloudSize - 5; ++i)
    {
        float depth1 = range_vec[i];
        float depth2 = range_vec[i + 1];
        if (depth1 - depth2 > 0.03 * depth2)
        {
            cloudNeighborPicked[i - 5] = 1;
            cloudNeighborPicked[i - 4] = 1;
            cloudNeighborPicked[i - 3] = 1;
            cloudNeighborPicked[i - 2] = 1;
            cloudNeighborPicked[i - 1] = 1;
            cloudNeighborPicked[i] = 1;
            BadcloudMarked[i - 5] = 1;
            BadcloudMarked[i - 4] = 1;
            BadcloudMarked[i - 3] = 1;
            BadcloudMarked[i - 2] = 1;
            BadcloudMarked[i - 1] = 1;
            BadcloudMarked[i] = 1;
        }
        else if (depth2 - depth1 > 0.03 * depth1)
        {
            cloudNeighborPicked[i + 1] = 1;
            cloudNeighborPicked[i + 2] = 1;
            cloudNeighborPicked[i + 3] = 1;
            cloudNeighborPicked[i + 4] = 1;
            cloudNeighborPicked[i + 5] = 1;
            cloudNeighborPicked[i + 6] = 1;

            BadcloudMarked[i + 1] = 1;
            BadcloudMarked[i + 2] = 1;
            BadcloudMarked[i + 3] = 1;
            BadcloudMarked[i + 4] = 1;
            BadcloudMarked[i + 5] = 1;
            BadcloudMarked[i + 6] = 1;
        }
        float diff1 = std::abs(range_vec[i - 1] - range_vec[i]);
        float diff2 = std::abs(range_vec[i + 1] - range_vec[i]);
        if (diff1 > 0.015 * range_vec[i] && diff2 > 0.015 * range_vec[i])
        {
            cloudNeighborPicked[i] = 1;
            BadcloudMarked[i] = 1;
        }
    }

    TicToc t_pts;
    pcl::PointCloud<PointType> cornerPointsSharp;
    pcl::PointCloud<PointType> cornerPointsLessSharp;
    pcl::PointCloud<PointType> surfPointsFlat;
    pcl::PointCloud<PointType> surfPointsLessFlat;
    float t_q_sort = 0;
    for (int i = 0; i < N_SCANS; i++)
    {
        if (scanEndInd[i] - scanStartInd[i] < 6) continue;
        pcl::PointCloud<PointType>::Ptr surfPointsLessFlatScan(new pcl::PointCloud<PointType>);
        for (int j = 0; j < 6; j++)
        {
            int sp = scanStartInd[i] + (scanEndInd[i] - scanStartInd[i]) * j / 6;
            int ep = scanStartInd[i] + (scanEndInd[i] - scanStartInd[i]) * (j + 1) / 6 - 1;
            TicToc t_tmp;
            std::sort(cloudSortInd + sp, cloudSortInd + ep + 1, comp);
            t_q_sort += t_tmp.toc();
            int largestPickedNum = 0;
            for (int k = ep; k >= sp; k--)
            {
                int ind = cloudSortInd[k];
                if (cloudNeighborPicked[ind] == 0 && cloudCurvature[ind] > 0.1)
                {
                    largestPickedNum++;
                    if (largestPickedNum <= 2)
                    {
                        cloudLabel[ind] = 2;
                        cornerPointsSharp.push_back(laserCloud->points[ind]);
                        cornerPointsLessSharp.push_back(laserCloud->points[ind]);
                    }
                    else if (largestPickedNum <= 20)
                    {
                        cloudLabel[ind] = 1;
                        cornerPointsLessSharp.push_back(laserCloud->points[ind]);
                    }
                    else break;
                    cloudNeighborPicked[ind] = 1;
                    for (int l = 1; l <= 5; l++)
                    {
                        float diffX = laserCloud->points[ind + l].x - laserCloud->points[ind + l - 1].x;
                        float diffY = laserCloud->points[ind + l].y - laserCloud->points[ind + l - 1].y;
                        float diffZ = laserCloud->points[ind + l].z - laserCloud->points[ind + l - 1].z;
                        if (diffX * diffX + diffY * diffY + diffZ * diffZ > 0.05) break;
                        cloudNeighborPicked[ind + l] = 1;
                    }
                    for (int l = -1; l >= -5; l--)
                    {
                        float diffX = laserCloud->points[ind + l].x - laserCloud->points[ind + l + 1].x;
                        float diffY = laserCloud->points[ind + l].y - laserCloud->points[ind + l + 1].y;
                        float diffZ = laserCloud->points[ind + l].z - laserCloud->points[ind + l + 1].z;
                        if (diffX * diffX + diffY * diffY + diffZ * diffZ > 0.05) break;
                        cloudNeighborPicked[ind + l] = 1;
                    }
                }
            }
            int smallestPickedNum = 0;
            for (int k = sp; k <= ep; k++)
            {
                int ind = cloudSortInd[k];
                if (cloudNeighborPicked[ind] == 0 && cloudCurvature[ind] < 0.1)
                {
                    cloudLabel[ind] = -1;
                    surfPointsFlat.push_back(laserCloud->points[ind]);
                    smallestPickedNum++;
                    if (smallestPickedNum >= 4) break;
                    cloudNeighborPicked[ind] = 1;
                    for (int l = 1; l <= 5; l++)
                    {
                        float diffX = laserCloud->points[ind + l].x - laserCloud->points[ind + l - 1].x;
                        float diffY = laserCloud->points[ind + l].y - laserCloud->points[ind + l - 1].y;
                        float diffZ = laserCloud->points[ind + l].z - laserCloud->points[ind + l - 1].z;
                        if (diffX * diffX + diffY * diffY + diffZ * diffZ > 0.05) break;
                        cloudNeighborPicked[ind + l] = 1;
                    }
                    for (int l = -1; l >= -5; l--)
                    {
                        float diffX = laserCloud->points[ind + l].x - laserCloud->points[ind + l + 1].x;
                        float diffY = laserCloud->points[ind + l].y - laserCloud->points[ind + l + 1].y;
                        float diffZ = laserCloud->points[ind + l].z - laserCloud->points[ind + l + 1].z;
                        if (diffX * diffX + diffY * diffY + diffZ * diffZ > 0.05) break;
                        cloudNeighborPicked[ind + l] = 1;
                    }
                }
            }
            for (int k = sp; k <= ep; k++)
            {
                if (cloudLabel[k] <= 0 && BadcloudMarked[k] == 0)
                {
                    surfPointsLessFlatScan->push_back(laserCloud->points[k]);
                }
            }
        }
        pcl::PointCloud<PointType> surfPointsLessFlatScanDS;
        pcl::VoxelGrid<PointType> downSizeFilter;
        downSizeFilter.setInputCloud(surfPointsLessFlatScan);
        downSizeFilter.setLeafSize(0.4, 0.4, 0.4);
        downSizeFilter.filter(surfPointsLessFlatScanDS);
        surfPointsLessFlat += surfPointsLessFlatScanDS;
    }
    sensor_msgs::PointCloud2 laserCloudOutMsg;
    pcl::toROSMsg(*laserCloud, laserCloudOutMsg);
    laserCloudOutMsg.header.stamp = laserCloudMsg->header.stamp;
    laserCloudOutMsg.header.frame_id = "base_link";
    pubLaserCloud.publish(laserCloudOutMsg);

    sensor_msgs::PointCloud2 cornerPointsSharpMsg;
    pcl::toROSMsg(cornerPointsSharp, cornerPointsSharpMsg);
    cornerPointsSharpMsg.header.stamp = laserCloudMsg->header.stamp;
    cornerPointsSharpMsg.header.frame_id = "base_link";
    pubCornerPointsSharp.publish(cornerPointsSharpMsg);

    sensor_msgs::PointCloud2 cornerPointsLessSharpMsg;
    pcl::toROSMsg(cornerPointsLessSharp, cornerPointsLessSharpMsg);
    cornerPointsLessSharpMsg.header.stamp = laserCloudMsg->header.stamp;
    cornerPointsLessSharpMsg.header.frame_id = "base_link";
    pubCornerPointsLessSharp.publish(cornerPointsLessSharpMsg);

    sensor_msgs::PointCloud2 surfPointsFlat2;
    pcl::toROSMsg(surfPointsFlat, surfPointsFlat2);
    surfPointsFlat2.header.stamp = laserCloudMsg->header.stamp;
    surfPointsFlat2.header.frame_id = "base_link";
    pubSurfPointsFlat.publish(surfPointsFlat2);

    sensor_msgs::PointCloud2 surfPointsLessFlat2;
    pcl::toROSMsg(surfPointsLessFlat, surfPointsLessFlat2);
    surfPointsLessFlat2.header.stamp = laserCloudMsg->header.stamp;
    surfPointsLessFlat2.header.frame_id = "base_link";
    pubSurfPointsLessFlat.publish(surfPointsLessFlat2);

    q_odom_last = q_odom_curr;
    t_odom_last = t_odom_curr;
    // pub each scam
    if (PUB_EACH_LINE)
    {
        for (int i = 0; i < N_SCANS; i++)
        {
            sensor_msgs::PointCloud2 scanMsg;
            pcl::toROSMsg(laserCloudScans[i], scanMsg);
            scanMsg.header.stamp = laserCloudMsg->header.stamp;
            scanMsg.header.frame_id = "base_link";
            pubEachScan[i].publish(scanMsg);
        }
    }
    //   printf("scan registration time %f ms *************\n", t_whole.toc());
    if (t_whole.toc() > 100) ROS_WARN("scan registration process over 100ms");
}

int main(int argc, char **argv)
{
    ros::init(argc, argv, "scanRegistration");
    ros::NodeHandle nh;
    nh.param<int>("scan_line", N_SCANS, 16);
    nh.param<double>("minimum_range", MINIMUM_RANGE, 0.1);
    printf("scan line number %d \n", N_SCANS);
    if (N_SCANS != 16 && N_SCANS != 32 && N_SCANS != 64)
    {
        printf("only support velodyne with 16, 32 or 64 scan line!");
        return 0;
    }
    ros::Subscriber subLaserCloud = nh.subscribe<sensor_msgs::PointCloud2>("rslidar_points", 100, laserCloudHandler);
    ros::Subscriber subOdom = nh.subscribe<nav_msgs::Odometry>("robot_fusion_odom", 100, odomHandler);
    pubLaserCloud = nh.advertise<sensor_msgs::PointCloud2>("velodyne_cloud_2", 100);
    pubCornerPointsSharp = nh.advertise<sensor_msgs::PointCloud2>("laser_cloud_sharp", 100);
    pubCornerPointsLessSharp = nh.advertise<sensor_msgs::PointCloud2>("laser_cloud_less_sharp", 100);
    pubSurfPointsFlat = nh.advertise<sensor_msgs::PointCloud2>("laser_cloud_flat", 100);
    pubSurfPointsLessFlat = nh.advertise<sensor_msgs::PointCloud2>("laser_cloud_less_flat", 100);
    pubRemovePoints = nh.advertise<sensor_msgs::PointCloud2>("laser_remove_points", 100);
    if (PUB_EACH_LINE)
    {
        for (int i = 0; i < N_SCANS; i++)
        {
            ros::Publisher tmp = nh.advertise<sensor_msgs::PointCloud2>("laser_scanid_" + std::to_string(i), 100);
            pubEachScan.push_back(tmp);
        }
    }
    ros::spin();
    return 0;
}
